今天會做到在專案中使用 Rust 做出推論 底下先列出 tasks
import torch
import numpy as np
# 載入權重
state_dict = torch.load("model.pt", map_location="cpu")
# 轉存為 npz
npz_weights = {}
for k, v in state_dict.items():
npz_weights[k] = v.numpy()
np.savez("model.npz", **npz_weights)
ex :
產出一個 model.npz,裡面包含:
['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias']
晚點會放上 實作細節